package org.dromara.common.websocket.handler;

import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.unfbx.chatgpt.OpenAiStreamClient;
import com.unfbx.chatgpt.entity.chat.Message;
import org.dromara.common.core.domain.model.LoginUser;
import org.dromara.common.core.utils.SpringUtils;
import org.dromara.common.websocket.config.LocalCache;
import org.dromara.common.websocket.holder.WebSocketSessionHolder;
import org.dromara.common.websocket.listener.WebSocketEventListener;
import org.dromara.common.websocket.service.TextReviewService;
import org.dromara.common.websocket.utils.WebSocketUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.AbstractWebSocketHandler;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

import static org.dromara.common.websocket.constant.WebSocketConstants.LOGIN_USER_KEY;

/**
 * WebSocketHandler 实现类
 *
 * @author zendwang
 */
@Slf4j
public class PlusWebSocketHandler extends AbstractWebSocketHandler {

    /**
     * 上下文长度
     */
    private final int contextSize = 10;

    /**
     * 审核状态 1 代表合法
     */
    static final String conclusionType = "1";

    private OpenAiStreamClient openAiStreamClient;

    private TextReviewService textReviewService;

    /**
     * 连接成功后
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
       // LoginUser loginUser = (LoginUser) session.getAttributes().get(LOGIN_USER_KEY);
        WebSocketSessionHolder.addSession(session.getId(), session);
        log.info("[connect] sessionId: {},userId:{},userName:{}", session.getId());
    }

    /**
     * 处理发送来的文本消息
     *
     * @param session
     * @param message
     * @throws Exception
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
        // 判断文本是否合规
//        textReviewService = (TextReviewService) SpringUtils.context().getBean("textReviewService");
//        String type = textReviewService.textReview(message.getPayload());
//        if(!conclusionType.equals(type)){
//            HashMap<Object, Object> msgMap = new HashMap<>();
//            msgMap.put("content","文本不合规,请修改后重试");
//            String jsonStr = JSONUtil.toJsonStr(msgMap);
//            WebSocketUtils.sendMessage(session, jsonStr);
//            WebSocketUtils.sendMessage(session,"[DONE]");
//            return;
//        }
      //  LoginUser loginUser = (LoginUser) session.getAttributes().get(LOGIN_USER_KEY);
        log.info("PlusWebSocketHandler, 连接：" + session.getId() + "，已收到消息:" + message.getPayload() );
        WebSocketEventListener eventSourceListener = new WebSocketEventListener(session);
        String messageContext = (String) LocalCache.CACHE.get(session.getId());
        List<Message> messages = new ArrayList<>();
        if (StrUtil.isNotBlank(messageContext)) {
            messages = JSONUtil.toList(messageContext, Message.class);
            if (messages.size() >= contextSize) {
                messages = messages.subList(1, contextSize);
            }
            Message currentMessage = Message.builder().content(message.getPayload()).role(Message.Role.USER).build();
            messages.add(currentMessage);
        } else {
            Message currentMessage = Message.builder().content(message.getPayload()).role(Message.Role.USER).build();
            messages.add(currentMessage);
        }
        openAiStreamClient = (OpenAiStreamClient) SpringUtils.context().getBean("openAiStreamClient");
        openAiStreamClient.streamChatCompletion(messages, eventSourceListener);
        LocalCache.CACHE.put(session.getId(), JSONUtil.toJsonStr(messages), LocalCache.TIMEOUT);
    }

    @Override
    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) throws Exception {
        super.handleBinaryMessage(session, message);
    }

    /**
     * 心跳监测的回复
     *
     * @param session
     * @param message
     * @throws Exception
     */
    @Override
    protected void handlePongMessage(WebSocketSession session, PongMessage message) throws Exception {
        WebSocketUtils.sendPongMessage(session);
    }

    /**
     * 连接出错时
     *
     * @param session
     * @param exception
     * @throws Exception
     */
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.error("[transport error] sessionId: {} , exception:{}", session.getId(), exception.getMessage());
    }

    /**
     * 连接关闭后
     *
     * @param session
     * @param status
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
      //  LoginUser loginUser = (LoginUser) session.getAttributes().get(LOGIN_USER_KEY);
        WebSocketSessionHolder.removeSession(session.getId());
       // log.info("[disconnect] sessionId: {},userId:{},userType:{}", session.getId(), loginUser.getUserId(), loginUser.getUserType());
    }

    /**
     * 是否支持分片消息
     *
     * @return
     */
    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

}
